from debugtools import *
from codepickle import *
import multiprocessing, multiprocessing.connection, threading, logging
import os, sys, zlib, cPickle, time, traceback, gc, socket, base64, math, binascii


SOCKETTIMES = [.25, .5, 1.0, 2.0] # if connection try number i fails, try again after SOCKETTIMES[i] seconds
BUFSIZE = 2048


class Connection(object):
    """
    Handles chunking and compression of data.
    
    To minimise data transfers between machines, we can use data compression,
    which this Connection handles automatically.
    """
    def __init__(self, conn, chunked=True, compressed=False):
        self.conn = conn
        self.chunked = chunked
        self.compressed = compressed
        self.BUFSIZE = BUFSIZE
        
    def send(self, obj):
        s = cPickle.dumps(obj, -1)
        if self.compressed:
            s = zlib.compress(s)
        if self.chunked:
            l = int(math.ceil(float(len(s))/self.BUFSIZE))
            # len(s) is a multiple of BUFSIZE, padding right with spaces
            s = s.ljust(l*self.BUFSIZE)
            l = "%08d" % l
            self.conn.sendall(l)
            self.conn.sendall(s)
        else:
            self.conn.sendall(s)
            
    def recv(self):
        if self.chunked:
            # Gets the first 8 bytes to retrieve the number of packets.
            l = ""
            n = 8
            while n > 0:
                l += self.conn.recv(n)
                n -= len(l)
            # BUG: sometimes l is filled with spaces??? setting l=1 in this case
            # (not a terrible solution)
            try:
                l = int(l)
            except:
                log.warn("transfer error, the paquet size was empty")
                l = 1
            
            length = l*self.BUFSIZE
            s = ""
            # Ensures that all data have been received
            while len(s) < length:
                data = self.conn.recv(self.BUFSIZE)
                s += data
        else:
            s = self.conn.recv()
        # unpads spaces on the right
        s = s.rstrip()
        if self.compressed:
            s = zlib.decompress(s)
        return cPickle.loads(s)
    
    def close(self):
        if self.conn is not None:
            r = self.conn.close()
            self.conn = None




def accept(address):
    """
    Accepts a connection and returns a connection object.
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    for i in xrange(5):
        try:
#            log.debug("trying to bind the socket...")
            s.bind(address)
            s.listen(5)
#            log.debug("the socket is now listening")
            break
        except:
            if i<4:
                t = SOCKETTIMES[i]
                log.debug("unable to bind the socket, trying again in %.2f seconds..." % t)
                time.sleep(t)
            else:
                log.debug("unable to bind the socket")
                raise Exception("unable to bind the socket")
    try:
        conn, addr = s.accept()
    except:
        raise Exception("unable to accept incoming connections")
    conn = Connection(conn)
    
    # The client can send its id at each connection, otherwise it can be retrieved
    # from the socket.accept() function.
#    clientid = conn.recv()
#    if clientid is None:
#        clientid = addr[0]
#    log.debug("client address: %s" % clientid)
    
    return conn, addr[0]

def connect(address):
    """
    Connects to a server and returns a Connection object.
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    for i in xrange(5):
        try:
#            log.debug("client: trying to connect to %s..." % address[0])
            s.connect(address)
#            log.debug("client: connection ok")
            break
        except:
            if i<4:
                t = SOCKETTIMES[i]
                log.debug("client: unable to connect, trying again in %.2f seconds..." % t)
                time.sleep(t)
            else:
                msg = "unable to connect to '%s' on port %d" % address
                log.warn(msg)
                raise Exception(msg)
                return None
    conn = Connection(s)
    # sends the client id to the server
#    conn.send(clientid)
    time.sleep(.05) # waits a bit for the server to be ready to receive
    return conn